package com.by.rpc.factory;

import com.by.log.Logger;
import com.by.log.LoggerFactory;
import com.by.rpc.netty.RpcClient;
import com.by.rpc.pool.RpcConnectionPool;
import lombok.Getter;

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * RPC客户端工厂，用于创建和管理RPC客户端实例
 */
public class RpcClientFactory {
    private static final Logger logger = LoggerFactory.getLogger(RpcClientFactory.class);

    @Getter
    private static final RpcClientFactory instance = new RpcClientFactory();

    private final ConcurrentHashMap<String, RpcConnectionPool> connectionPools = new ConcurrentHashMap<>();

    // 新增独立的线程池，用于处理阻塞的连接获取操作
    private final ExecutorService connectionExecutor = new ThreadPoolExecutor(
            1, 10, 60L, TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(),
            new ThreadFactory() {
                private final AtomicInteger counter = new AtomicInteger(0);
                @Override
                public Thread newThread(Runnable r) {
                    return new Thread(r, "RpcClientFactory-Connector-" + counter.incrementAndGet());
                }
            });

    private RpcClientFactory() {
        logger.info("Initializing RpcClientFactory");
    }

    /**
     * 异步获取RPC客户端
     */
    public CompletableFuture<RpcClient> getClientAsync(String host, int port) {
        String key = host + ":" + port;
        logger.debug("Async getting client for {}:{}", host, port);

        RpcConnectionPool pool = connectionPools.computeIfAbsent(key, k -> {
            logger.info("Creating new connection pool for {}", k);
            return new RpcConnectionPool(host, port, 10, 2, 60000);
        });

        return pool.getConnectionAsync();
    }

    /**
     * 归还RPC客户端
     */
    public void returnClient(String host, int port, RpcClient client) {
        try {
            String key = host + ":" + port;
            RpcConnectionPool pool = connectionPools.get(key);
            if (pool != null) {
                pool.returnConnection(client);
            } else {
                logger.warn("No connection pool found for {}, closing client directly", key);
                client.close();
            }
        } catch (Exception e) {
            logger.error("Error returning client for {}:{}", host, port, e);
        }
    }

    /**
     * 关闭所有连接池
     */
    public void shutdown() {
        logger.info("Shutting down all connection pools");
        for (RpcConnectionPool pool : connectionPools.values()) {
            pool.close();
        }
        connectionPools.clear();
        connectionExecutor.shutdown();
        try {
            if (!connectionExecutor.awaitTermination(5, TimeUnit.SECONDS)) {
                connectionExecutor.shutdownNow();
            }
        } catch (InterruptedException e) {
            connectionExecutor.shutdownNow();
            Thread.currentThread().interrupt();
        }
        logger.info("All connection pools shut down");
    }
}